import os
import json
import cv2
import numpy as np


# 加载 JSON 文件
def load_json(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data


# 提取右侧和左侧关键点
def extract_points(data):
    right_labels = ['Rlong1', 'Rlong2', 'Rass1', 'Rass2', 'Rass3', 'Rass4']
    left_labels = ['Llong1', 'Llong2', 'Lass1', 'Lass2', 'Lass3', 'Lass4']

    right_points = []
    left_points = []

    for shape in data['shapes']:
        label = shape['label']
        point = shape['points'][0]  # 获取点坐标
        if label in right_labels:
            right_points.append(tuple(point))
        elif label in left_labels:
            left_points.append(tuple(point))

    return np.array(right_points, dtype=np.float32), np.array(left_points, dtype=np.float32)


# 检查是否有14个点, 未置换时无需拟合
def check_points(data):
    return len(data['shapes']) == 14


# 拟合椭圆并绘制
def fit_and_draw_ellipse(points, canvas, color):
    if len(points) >= 5:  # 椭圆拟合需要至少 5 个点
        ellipse = cv2.fitEllipse(points)
        cv2.ellipse(canvas, ellipse, color, 2)
    else:
        print("Not enough points to fit an ellipse.")
    return canvas


# 绘制关键点
def draw_points(points, canvas, color):
    for point in points:
        cv2.circle(canvas, (int(point[0]), int(point[1])), 5, color, -1)
    return canvas


# 保存图片到指定路径
def save_image(image, output_path):
    os.makedirs(os.path.dirname(output_path), exist_ok=True)  # 确保路径存在
    cv2.imwrite(output_path, image)


# 处理每个JSON文件，拟合椭圆并保存结果
def process_json_file(json_path, output_folder):
    data = load_json(json_path)

    # 检查文件是否标注了14个点
    if not check_points(data):
        print(f"无需拟合：{json_path}")
        return

    # 获取图像的尺寸
    image_width = data["imageWidth"]
    image_height = data["imageHeight"]

    # 提取右侧和左侧的关键点
    right_points, left_points = extract_points(data)

    # 创建一个空白画布
    canvas = np.ones((image_height, image_width, 3), dtype=np.uint8) * 255

    # 绘制原始关键点
    canvas = draw_points(right_points, canvas, (0, 255, 0))  # 绿色点表示右侧
    canvas = draw_points(left_points, canvas, (255, 0, 0))  # 红色点表示左侧

    # 对右侧和左侧关键点进行椭圆拟合并绘制
    canvas = fit_and_draw_ellipse(right_points, canvas, (0, 255, 0))  # 绿色椭圆表示右侧
    canvas = fit_and_draw_ellipse(left_points, canvas, (255, 0, 0))  # 红色椭圆表示左侧

    # 获取文件名并去掉后缀
    filename = os.path.basename(json_path).replace('.json', '.png')
    output_path = os.path.join(output_folder, filename)

    # 保存图片
    save_image(canvas, output_path)


# 遍历所有文件夹并处理每个JSON文件
def process_all_folders(data_folder, output_base_folder):
    for root, dirs, files in os.walk(data_folder):
        for file in files:
            if file.endswith('.json'):
                json_path = os.path.join(root, file)
                # 构建输出文件夹路径，替换掉 'DataMining_Dataset' 为 'result/ellipse'
                relative_path = os.path.relpath(root, data_folder)
                output_folder = os.path.join(output_base_folder, relative_path)

                # 处理该 JSON 文件并保存结果
                process_json_file(json_path, output_folder)


# 主函数
def main():
    data_folder = 'image/DataMining_Dataset'  # 输入数据集的根文件夹
    output_base_folder = 'result/ellipse'  # 输出结果的根文件夹

    # 处理所有文件夹下的 JSON 文件
    process_all_folders(data_folder, output_base_folder)


if __name__ == '__main__':
    main()
